#!/usr/bin/env Rscript
library(optparse)
library(leafcutter)

arguments <- parse_args(OptionParser(usage = "%prog [options] counts_file groups_file cluster_significance_file", description="LeafCutter differential splicing plotting tool. Required inputs:\n <counts_file>: Intron usage counts file. Must be .txt or .txt.gz, output from clustering pipeline.\n <groups_file>: Two column file: 1. sample names (must match column names in counts_file), 2. groups (currently only two groups, i.e. pairwise, supported. Some samples in counts_file can be missing from this file, in which case they will not be included in the analysis. \n<cluster_significance_file> generated by leafcutter_ds.R",option_list=list(
  make_option(c("-o","--output"), default = "ds_plots.pdf", help="The output file  [default %default]"),
  make_option(c("-f","--plot_FDR"), default=0.01, help="FDR (in 0-1) below which to make plots  [default %default]"), 
  make_option(c("-m","--max_plots"), default=30, help="Maximum number of plots to create [default %default] "), 
  make_option(c("-e","--exon_file"), default=NULL, help="File defining known exons, example in data/gencode19_exons.txt.gz. Columns should be chr, start, end, strand, gene_name."))), 
  positional_arguments = 3)

opt=arguments$opt
counts_file=arguments$args[1]
groups_file=arguments$args[2]
cluster_significance_file=arguments$args[3]

cat("Loading counts from",counts_file,"\n")
if (!file.exists(counts_file)) stop("File ",counts_file," does not exist")
counts=read.table(counts_file, header=T, check.names=F, stringsAsFactors = F)

cat("Loading metadata from",groups_file,"\n")
if (!file.exists(groups_file)) stop("File ",groups_file," does not exist")
meta=read.table(groups_file, header=F, stringsAsFactors = F)
colnames(meta)=c("sample","group")

exon_table=if (!is.null(opt$exon_file)) {
  cat("Loading exons from",opt$exon_file,"\n")
  if (!file.exists(opt$exon_file)) stop("File ",opt$exon_file," does not exist")
  read.table(opt$exon_file, header=T, stringsAsFactors = F)
} else {
  cat("No exon_file provided.\n")
  NULL
}

counts=counts[,meta$sample]

meta$group=as.factor(meta$group)
group_names=levels(meta$group)

stopifnot(length(group_names)==2)

if (!file.exists(cluster_significance_file)) stop("File ",cluster_significance_file," does not exist")
cluster_table=read.table(cluster_significance_file, header=T, sep="\t", stringsAsFactors = F)

qvalues=leafcutter:::bh(cluster_table$p)
top_clus=cluster_table$cluster[ which(qvalues < opt$plot_FDR) ]

if (length(top_clus) > opt$max_plots) 
  top_clus=cluster_table$cluster[ order(qvalues)[seq_len(opt$max_plots)] ]

cat("Saving",length(top_clus),"plots to",opt$output,"\n")
introns=get_intron_meta(rownames(counts))
cluster_ids=paste( add_chr(introns$chr),introns$clu,sep = ":")

pdf(opt$output, width=8, height=8)
for (clu in top_clus) {
  print(clu)
  y=t(counts[ cluster_ids==clu, ])
  make_differential_splicing_plot(y, meta$group, exons_table=exon_table)
}
dev.off()

cat("All done, exiting\n")
